package com.wl.cloud.datax.common.utils;

import cn.hutool.extra.ssh.JschUtil;
import com.jcraft.jsch.Session;
import com.wl.cloud.core.exception.CustomException;
import com.wl.cloud.datax.common.config.KafkaProperties;
import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * @Author: panzhichun
 * @Date: 2024/2/22 13:50
 * @Description:
 * @Version: 1.0
 */
@Slf4j
public class CustomKafkaUserUtils {

    // 增加kafka用户
    public static void addKafkaUser(KafkaProperties kafkaProperties, String kafkaUser, String kafkaPass) {
        // 查询zookeeper中原本是否存在该用户
        List<String> allKafkaUser;
        allKafkaUser = getAllKafkaUser(kafkaProperties.getZookeeperAddr());
        if (allKafkaUser.contains(kafkaUser)) {
            throw new CustomException("用户[" + kafkaUser + "]已存在");
        }
        // 获取session连接
        Session session = JschUtil.getSession(kafkaProperties.getSshIp(), kafkaProperties.getSshPort(), kafkaProperties.getSshUser(), kafkaProperties.getSshPass());
        // 拼接新增用户的脚本命令
        String cmd = kafkaProperties.getKafkaBinPath() + "kafka-configs.sh --zookeeper " + kafkaProperties.getZookeeperAddr() + " --alter --add-config '" + kafkaProperties.getSaslMechanism() + "=[password=%s]' --entity-type users --entity-name %s";
        cmd = String.format(cmd, kafkaPass, kafkaPass, kafkaUser);
        // 执行命令
        JschUtil.exec(session, cmd, StandardCharsets.UTF_8);
        // 关闭session
        JschUtil.close(session);
        // 查询zookeeper检验是否新增成功
        allKafkaUser = getAllKafkaUser(kafkaProperties.getZookeeperAddr());
        if (!allKafkaUser.contains(kafkaUser)) {
            throw new CustomException("用户[" + kafkaUser + "]新增失败");
        }
    }

    // 修改用户密码
    public static void updateKafkaUserPass(KafkaProperties kafkaProperties, String kafkaUser, String kafkaPass) {
        // 查询zookeeper检验是否新增成功
        List<String> allKafkaUser = getAllKafkaUser(kafkaProperties.getZookeeperAddr());
        if (!allKafkaUser.contains(kafkaUser)) {
            throw new CustomException("用户[" + kafkaUser + "]不存在");
        }
        // 获取session连接
        Session session = JschUtil.getSession(kafkaProperties.getSshIp(), kafkaProperties.getSshPort(), kafkaProperties.getSshUser(), kafkaProperties.getSshPass());
        // 拼接新增用户的脚本命令
        String cmd = kafkaProperties.getKafkaBinPath() + "kafka-configs.sh --zookeeper " + kafkaProperties.getZookeeperAddr() + " --alter --add-config '" + kafkaProperties.getSaslMechanism() + "=[password=%s]' --entity-type users --entity-name %s";
        cmd = String.format(cmd, kafkaPass, kafkaPass, kafkaUser);
        // 执行命令
        JschUtil.exec(session, cmd, StandardCharsets.UTF_8);
        // 关闭session
        JschUtil.close(session);
    }

    // 删除用户
    public static void deleteKafkaUser(KafkaProperties kafkaProperties, String kafkaUser) {
        try {
            // 获取session连接
            Session session = JschUtil.getSession(kafkaProperties.getSshIp(), kafkaProperties.getSshPort(), kafkaProperties.getSshUser(), kafkaProperties.getSshPass());
            // 拼接新增用户的脚本命令
            String cmd = kafkaProperties.getKafkaBinPath() + "kafka-configs.sh --zookeeper " + kafkaProperties.getZookeeperAddr() + " --alter --delete-config '" + kafkaProperties.getSaslMechanism() + "' --entity-type users --entity-name %s";
            cmd = String.format(cmd, kafkaUser);
            // 执行命令
            JschUtil.exec(session, cmd, StandardCharsets.UTF_8);
            // 关闭session
            JschUtil.close(session);
            //执行删除脚本后，zk节点还存在，删除zk中的用户节点
            ZooKeeper zookeeper = new ZooKeeper(kafkaProperties.getZookeeperAddr(), 300000, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    log.info("zookeeper删除用户节点，收到事件通知：{}", event.getState() );
                }
            });
            zookeeper.delete("/config/users/" + kafkaUser,-1);
            // 查询zookeeper校验是否删除成功
            List<String> allKafkaUser = getAllKafkaUser(kafkaProperties.getZookeeperAddr());
            if (allKafkaUser.contains(kafkaUser)) {
                throw new CustomException("用户[" + kafkaUser + "]删除失败");
            }
        } catch (Exception e) {
            log.warn("删除用户信息失败，失败信息：{}", e.getMessage());
        }
    }

    // 查询zookeeper中所有的kafka用户
    public static List<String> getAllKafkaUser(String zookeeperAddr) {
        List<String> list = new ArrayList<>();
        ZooKeeper zookeeper = null;
        try {
            zookeeper = new ZooKeeper(zookeeperAddr, 300000, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    log.info("查询zk用户列表信息，连接zk收到事件通知：" + event.getState() );
                }
            });
            list = zookeeper.getChildren("/config/users", new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    log.info("查询用户列表信息，收到事件通知：" + event.getState() );
                }
            });
        } catch (Exception e) {
            log.error("查询zk子节点信息失败，异常信息：{}", e.getMessage());
        }
        return list;
    }

}
